package com.navinfo.tripanalysis.service.impl;

import com.navinfo.tripanalysis.common.util.DateUtils;
import com.navinfo.tripanalysis.pojo.LoadDataParam;
import com.navinfo.tripanalysis.util.CommonUtils;
import lombok.Data;
import org.apache.commons.lang3.StringUtils;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.api.java.function.Function;
import org.apache.spark.api.java.function.Function2;
import org.apache.spark.api.java.function.PairFlatMapFunction;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import scala.Tuple2;

import java.util.Iterator;
import java.util.List;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * 从hive表加载数据公共代码
 * @author 沈东生
 */
@Data
public abstract class AbstractLoadDataHive<T> {
    protected final Logger logger = LoggerFactory.getLogger(getClass());
    String tablename;
    String dbname;


    public JavaPairRDD<Long, List<T>> load(SparkSession spark, JavaSparkContext jsc, LoadDataParam param) {
        String dayStr = DateUtils.format(param.getDay(), DateUtils.DateFormat.YYYYMMDD);
        StringBuilder sb = new StringBuilder();
        sb.append("select * ");
        sb.append("\n");
        sb.append(" from ");
        if (StringUtils.isNotEmpty(dbname)) {
            sb.append(dbname).append(".");
        }
        sb.append(tablename);
        sb.append("\n");
        sb.append(" where ");
        sb.append(String.format(" day ='%s'", dayStr));

        List<Long> tids = param.getTids();
        if (CommonUtils.isCollectionNotEmpty(tids)) {
            sb.append(String.format(" and tid in(%s) ", CommonUtils.mkString(tids, ",")));
        }
        String sql = sb.toString();
        logger.error("sql:-----------------\n{}", sql);
        return spark.sql(sql).toJavaRDD().mapPartitionsToPair((PairFlatMapFunction<Iterator<Row>, Long, T>) rowIterator -> new Iterator<Tuple2<Long, T>>() {
            @Override
            public boolean hasNext() {
                return rowIterator.hasNext();
            }

            @Override
            public Tuple2<Long, T> next() {
                Row row = rowIterator.next();
                return new Tuple2<>((Long) row.getAs("tid"), fromRow(row));
            }
        }).combineByKey(
                (Function<T, List<T>>) point -> Stream.of(point).collect(Collectors.toList()),
                (Function2<List<T>, T, List<T>>) (points, p) ->{
                    points.add(p);
                    return points;
                },
                (Function2<List<T>, List<T>, List<T>>) (p1, p2) ->{
                    p1.addAll(p2);
                    return p1;
                });
    }

    /**
     * 将row转为具体的类型
     * @param row
     * @return
     */
    protected abstract T fromRow(Row row);
}
